from typing import Any, Literal

import pandas as pd

from ..dataclass.geoProfiles import MaterialTable
from ..dataclass.tables import FieldMetadata, TableCollection, TableData
from ..dataclass.terminologies import GeoMaterialProps as GProps
from ..dataclass.terminologies import GeoResults as GRes
from ..dataclass.terminologies import Units
from ..pipeline.pipeData import StringAttributeSchema, UIAttributeSchema
from ..pipeline.pipeline import (
    PipeModule,
    Port,
    PortReference,
    module_decorator,
    status_manage,
)
from ..pipeline.portTypes import PortType, PortTypeHint


@status_manage
class MergeMaterialTable(PipeModule):
    """Merge the two material tables to a single material table."""

    def __init__(
        self,
        mname: str = "MergeMaterialTable",
        auto_run: bool = True,
        material_table1: PortTypeHint.MaterialTable | None = None,
        material_table2: PortTypeHint.MaterialTable | None = None,
        merge_by: Literal[
            GProps.MaterialID, GProps.LayerNumber, GProps.MaterialName
        ] = GProps.MaterialID,
    ) -> None:
        """Initialize a MergeMaterialTable object.

        Parameters
        ----------
        material_table1: PortTypeHint.MaterialTable | None, default: None
            The main material table.

        material_table2: PortTypeHint.MaterialTable | None, default: None
            The another material table.

        merge_by: Literal[GProps.MaterialID, GProps.LayerNumber, GProps.MaterialName], default: GProps.MaterialID
            The column name to merge the material tables by.

        Note
        ----
        The matearial_table2 will be merged into the material_table1 which means the output material_table will have
        the same rows as the material_table1.
        """
        super().__init__(mname=mname, auto_run=auto_run)

        self._ports_in = {
            "InputMaterialTable1": Port(
                ptype=PortType.MaterialTable,
                data=material_table1,
                pdoc="The main material table.",
            ),
            "InputMaterialTable2": Port(
                ptype=PortType.MaterialTable,
                data=material_table2,
                pdoc="The another material table.",
            ),
        }
        self._ports_out = {
            "OutputMaterialTable": Port(
                ptype=PortType.MaterialTable, pdoc="The output material table."
            )
        }

        self.merge_by = merge_by

    @property
    def InputMaterialTable1(self) -> None:
        raise AttributeError("Property 'InputMaterialTable1' is write-only.")

    @InputMaterialTable1.setter
    def InputMaterialTable1(self, value: PortTypeHint.MaterialTable) -> None:
        self["InputMaterialTable1"] = value

    @property
    def InputMaterialTable2(self) -> None:
        raise AttributeError("Property 'InputMaterialTable2' is write-only.")

    @InputMaterialTable2.setter
    def InputMaterialTable2(self, value: PortTypeHint.MaterialTable) -> None:
        self["InputMaterialTable2"] = value

    @property
    def OutputMaterialTable(self) -> PortTypeHint.MaterialTable | None:
        return self["OutputMaterialTable"]

    def set_cal_params(self, reset: bool = True) -> None:
        return None

    def execute(self) -> PortTypeHint.MaterialTable | None:
        material_table1: PortTypeHint.MaterialTable | None = self._ports_in[
            "InputMaterialTable1"
        ].data
        material_table2: PortTypeHint.MaterialTable | None = self._ports_in[
            "InputMaterialTable2"
        ].data

        if material_table1 is None or material_table2 is None:
            self._ports_out["OutputMaterialTable"].data = None
            return None

        # Merge the two material tables
        material_table1.merge(material_table2, on=self.merge_by)

        self._ports_out["OutputMaterialTable"].data = material_table1
        return material_table1


@status_manage
class MergeBoresWithMaterialTable(PipeModule):
    """Merge Bores (MultiProfile1D) with MaterialTable to a TableCollection."""

    def __init__(
        self,
        mname: str = "MergeBoresWithMaterialTable",
        auto_run: bool = True,
        bores: PortTypeHint.MultiProfile1D | None = None,
        material_table: PortTypeHint.MaterialTable | None = None,
        additional_fields: (
            list[
                Literal[
                    GRes.LayerBottomElevation,
                    GRes.LayerTopElevation,
                    GRes.LayerThickness,
                ]
            ]
            | None
        ) = None,
    ) -> None:
        """Initialize a MergeBoresWithMaterialTable object.

        Parameters
        ----------
        bores: PortTypeHint.MultiProfile1D | None, default: None
            The bores to merge.

        material_table: PortTypeHint.MaterialTable | None, default: None
            The material table to merge.

        additional_fields: list[Literal[GRes.LayerBottomElevation, GRes.LayerTopElevation, GRes.LayerThickness]], default: None
            The additional fields to put into the output material table.

        Note
        ----
        Each bore in the MultiProfile1D will merge with the material table.
        A MaterialTableCollection will be returned and the keys are bore numbers.
        """
        super().__init__(mname=mname, auto_run=auto_run)

        self._ports_in = {
            "InputBores": Port(
                ptype=PortType.MultiProfile1D,
                data=bores,
                pdoc="MultiProfile1D bores data.",
            ),
            "InputMaterialTable": Port(
                ptype=PortType.MaterialTable,
                data=material_table,
                pdoc="MaterialTable to merge.",
            ),
        }
        self._ports_out = {
            "OutputTables": Port(
                ptype=PortType.TableCollection, pdoc="The output table collection."
            )
        }

        self.additional_fields = additional_fields

    @property
    def InputBores(self) -> None:
        raise AttributeError("Property 'InputBores' is write-only.")

    @InputBores.setter
    def InputBores(self, value: PortTypeHint.MultiProfile1D) -> None:
        self["InputBores"] = value

    @property
    def InputMaterialTable(self) -> None:
        raise AttributeError("Property 'InputMaterialTable' is write-only.")

    @InputMaterialTable.setter
    def InputMaterialTable(self, value: PortTypeHint.MaterialTable) -> None:
        self["InputMaterialTable"] = value

    @property
    def OutputTables(self) -> PortTypeHint.TableCollection | None:
        return self["OutputTables"]

    def set_cal_params(self, reset: bool = True) -> None:
        return None

    def execute(self) -> PortTypeHint.TableCollection | None:
        bores: PortTypeHint.MultiProfile1D | None = self._ports_in["InputBores"].data
        material_table: PortTypeHint.MaterialTable | None = self._ports_in[
            "InputMaterialTable"
        ].data

        if bores is None or material_table is None:
            self._ports_out["OutputTables"].data = None
            return None

        material_table = material_table.convert_to_table_data()

        # Merge the material table with the bores
        table_collection = TableCollection()
        for bore in bores:
            layers_data = {
                GProps.MaterialID: bore.materials_id,
                GRes.LayerDepth: bore.layers_depth,
            }
            fields_meta = [
                FieldMetadata(
                    name=GProps.MaterialID,
                    title=GProps.MaterialID.title,
                    unit=Units.UNITLESS,
                ),
                FieldMetadata(
                    name=GRes.LayerDepth, title=GRes.LayerDepth.title, unit=Units.m
                ),
            ]
            if self.additional_fields:
                if GRes.LayerBottomElevation in self.additional_fields:
                    layers_data[GRes.LayerBottomElevation] = bore.layers_elevation
                    fields_meta.append(
                        FieldMetadata(
                            name=GRes.LayerBottomElevation,
                            title=GRes.LayerBottomElevation.title,
                            unit=Units.m,
                        )
                    )
                if GRes.LayerTopElevation in self.additional_fields:
                    layers_data[GRes.LayerTopElevation] = bore.layers_top_elevation
                    fields_meta.append(
                        FieldMetadata(
                            name=GRes.LayerTopElevation,
                            title=GRes.LayerTopElevation.title,
                            unit=Units.m,
                        )
                    )
                if GRes.LayerThickness in self.additional_fields:
                    layers_data[GRes.LayerThickness] = bore.layers_thickness
                    fields_meta.append(
                        FieldMetadata(
                            name=GRes.LayerThickness,
                            title=GRes.LayerThickness.title,
                            unit=Units.m,
                        )
                    )
            bore_table = pd.DataFrame(layers_data)
            # Convert it to a material table
            bore_table = TableData(data=bore_table, fields_meta=fields_meta)
            # Merge the bore table with the material table
            bore_table = bore_table.merge(material_table, on=GProps.MaterialID)
            bore_table.name = f"{bore.pnum}"
            bore_table.title = f"{bore.pnum}钻孔"
            table_collection.add_table(bore_table)

        self._ports_out["OutputTables"].data = table_collection
        return table_collection


@status_manage
class ConcatenateTables(PipeModule):
    """Concatenate several tables."""

    def __init__(
        self,
        mname: str = "ConcatenateTables",
        auto_run: bool = True,
        tables: PortTypeHint.TableCollection | None = None,
        add_table_name_columns: str | None = None,
        add_table_title_columns: str | None = None,
    ) -> None:
        """Initialize a ConcatenateTables object.

        Parameters
        ----------
        tables: PortTypeHint.TableCollection | None, default: None
            The tables to concatenate.

        add_table_name_columns: str | None, default: None
            If provided, adds a column with this name containing the table names.

        add_table_title_columns: str | None, default: None
            If provided, adds a column with this name containing the table titles.

        Note
        ----
        `NaN` value will be filled for tables that do not the columns in their original data.
        """
        super().__init__(mname=mname, auto_run=auto_run)

        self._ports_in = {
            "InputTables": Port(
                ptype=PortType.TableCollection,
                data=tables,
                pdoc="The tables to concatenate.",
            )
        }
        self._ports_out = {
            "OutputTable": Port(
                ptype=PortType.TableData, pdoc="The concatenated table."
            )
        }

        self.add_table_name_columns = add_table_name_columns
        self.add_table_title_columns = add_table_title_columns

    @property
    def InputTables(self) -> None:
        raise AttributeError("Property 'InputTables' is write-only.")

    @InputTables.setter
    def InputTables(self, value: PortTypeHint.TableCollection) -> None:
        self["InputTables"] = value

    @property
    def OutputTable(self) -> PortTypeHint.TableData | None:
        return self["OutputTable"]

    def set_cal_params(self, reset: bool = True) -> None:
        return None

    def execute(self) -> PortTypeHint.TableData | None:
        tables: PortTypeHint.TableCollection | None = self._ports_in["InputTables"].data
        if tables is None:
            self._ports_out["OutputTable"].data = None
            return None

        # If we need to add table name or title columns, process each table first
        if self.add_table_name_columns or self.add_table_title_columns:
            modified_tables = []
            for table in tables:
                # Create a copy of the table data
                table_copy = table.copy()

                # Add table name column if requested
                if self.add_table_name_columns:
                    table_name = (
                        table.name if hasattr(table, "name") and table.name else ""
                    )
                    table_copy[self.add_table_name_columns] = table_name

                # Add table title column if requested
                if self.add_table_title_columns:
                    table_title = (
                        table.title if hasattr(table, "title") and table.title else ""
                    )
                    table_copy[self.add_table_title_columns] = table_title

                modified_tables.append(table_copy)

            # Create a new TableCollection with modified tables
            modified_table_collection = TableCollection()
            for table in modified_tables:
                modified_table_collection.add_table(table)

            # Use the new TableData.concat method with modified tables
            output_table = TableData.concat(
                modified_table_collection, ignore_index=True
            )
        else:
            # Use the new TableData.concat method with original tables
            output_table = TableData.concat(tables, ignore_index=True)

        self._ports_out["OutputTable"].data = output_table
        return output_table


@module_decorator()
class MergeTables(PipeModule):
    """Merge multiple tables from a TableCollection or merge two TableData objects."""

    InputTables: PortReference[PortTypeHint.TableCollection | PortTypeHint.TableData]
    OutputTable: PortReference[PortTypeHint.TableData]

    _port_docs = {
        "InputTables": "The input tables to merge.",
        "OutputTable": "The merged table.",
    }

    def __init__(
        self,
        mname: str = "MergeTables",
        auto_run: bool = True,
        tables: PortTypeHint.TableCollection | PortTypeHint.TableData | None = None,
        join_key: str | None = None,
        tables_to_merge: list[str] | str | None = None,
        how: Literal["left", "right", "outer", "inner"] = "inner",
        suffixes: tuple[str, str] = ("_x", "_y"),
        sort: bool = False,
        copy: bool = True,
        indicator: bool = False,
        validate: (
            Literal["one_to_one", "one_to_many", "many_to_one", "many_to_many"] | None
        ) = None,
        ignore_index: bool = False,
    ) -> None:
        """Initialize a MergeTables object.

        Parameters
        ----------
        tables: PortTypeHint.TableCollection | PortTypeHint.TableData | None, default: None
            The input tables to merge. Can be a TableCollection containing multiple tables
            or a single TableData object.

        join_key: str | None, default: None
            The column name or field title to use for merging tables. If None, uses the first common
            column that exists in all tables. Tables without this column will be skipped.

        tables_to_merge: list[str] | str | None, default: None
            The names or titles of tables to include in the merge. If None, all tables
            in the collection will be considered. If a single string, converts to list.

        how: Literal["left", "right", "outer", "inner"], default: "inner"
            Type of merge to be performed.
            Use "inner" when you only want complete records (no missing data)
            Use "left" when you want to preserve all records from your main dataset
            Use "right" rarely used, when you want all records from the second dataset
            Use "outer" when you want a comprehensive view including all possible records

        suffixes: tuple[str, str], default: ("_x", "_y")
            Suffix to apply to overlapping column names.

        sort: bool, default: False
            Sort the result DataFrame by the join keys.

        copy: bool, default: True
            If False, avoid copying data when possible.

        indicator: bool, default: False
            Add a column to the output DataFrame indicating the source.

        validate: Literal["one_to_one", "one_to_many", "many_to_one", "many_to_many"] | None, default: None
            Validation check on the merge operation.
            Use "one_to_one" when:
                Merging master records (employee info + employee contact)
                Each ID should appear exactly once in both tables
                You want to ensure no data duplication
            Use "one_to_many" when:
                Merging dimension tables with fact tables
                One category has multiple items (department → employees)
                Left table is the "lookup" table
            Use "many_to_one" when:
                Merging fact tables with dimension tables
                Multiple records reference the same lookup value
                Right table is the "lookup" table
            Use "many_to_many" when:
                Complex relationships exist
                You expect duplicates in both tables
                Junction/bridge table scenarios
            Use None when:
                You're not sure about the relationship
                Performance is critical (validation adds overhead)
                You'll handle data quality issues separately

        ignore_index: bool, default: False
            If True, do not use the index values along the concatenation axis.
            The resulting axis will be labeled 0, 1, …, n - 1.

        Note
        ----
        If tables is a TableData object, it will be returned as-is.
        If tables is a TableCollection, tables will be merged sequentially using the join_key.
        Tables that don't contain the join_key will be skipped with a warning.
        """
        super().__init__(mname=mname, auto_run=auto_run)

        if tables is not None:
            self.InputTables = tables

        self.join_key = join_key
        self.tables_to_merge = tables_to_merge
        self.how = how
        self.suffixes = suffixes
        self.sort = sort
        self.copy = copy
        self.indicator = indicator
        self.validate = validate
        self.ignore_index = ignore_index

    def update_ui_schema(self, reset: bool = False) -> dict[str, UIAttributeSchema]:
        tables: PortTypeHint.TableCollection | PortTypeHint.TableData | None = (
            self._ports_in["InputTables"].data
        )

        join_key_schema = StringAttributeSchema(title="索引字段")

        if tables is not None:
            if isinstance(tables, TableCollection):
                all_join_keys = self._find_join_key(
                    tables.values(), return_all_join_keys=True
                )
                join_key_schema.selections = all_join_keys  # type: ignore
                selections_name = [
                    tables[0].name_to_title[column_name]
                    for column_name in all_join_keys
                ]
                join_key_schema.selections_name = selections_name

        return {"join_key": join_key_schema}

    def _find_join_key(
        self,
        tables_to_merge: list[PortTypeHint.TableData],
        return_all_join_keys: bool = False,
    ) -> str | list[str] | None:
        """Find the first column that exists in all tables.
        If `return_all_join_keys = True`, return all columns that exist in all tables.
        """
        if not tables_to_merge:
            return None

        # Get columns from the first table
        first_table = tables_to_merge[0]
        common_columns = set(first_table.columns)

        # Find intersection with all other tables
        for table in tables_to_merge[1:]:
            common_columns &= set(table.columns)

        if not common_columns:
            return None

        if return_all_join_keys:
            # Return all common columns as a comma-separated string (in order of appearance in first table)
            common_cols_ordered = [
                col for col in first_table.columns if col in common_columns
            ]
            return common_cols_ordered
        else:
            # Return the first common column (in order of appearance in first table)
            for col in first_table.columns:
                if col in common_columns:
                    return col
            return None

    def _resolve_column_name(
        self, table: PortTypeHint.TableData, column_identifier: str
    ) -> str | None:
        """Resolve field title to column name if necessary."""
        # First check if it's already a column name
        if column_identifier in table.columns:
            return column_identifier

        # Check if it's a field title
        if column_identifier in table.title_to_name:
            return table.title_to_name[column_identifier]

        return None

    def execute(self) -> PortTypeHint.TableData | None:
        tables_input: PortTypeHint.TableCollection | PortTypeHint.TableData | None = (
            self._ports_in["InputTables"].data
        )

        if tables_input is None:
            self._ports_out["OutputTable"].data = None
            return None

        # If input is a single TableData, return it as-is
        if isinstance(tables_input, TableData):
            self._ports_out["OutputTable"].data = tables_input
            return tables_input

        # Get tables to merge
        if self.tables_to_merge is None:
            # Use all tables in the collection
            tables_to_merge = list(tables_input.values())
        else:
            # Use only specified tables
            tables_to_merge = []
            # Convert single string to list
            table_identifiers = (
                [self.tables_to_merge]
                if isinstance(self.tables_to_merge, str)
                else self.tables_to_merge
            )
            for table_identifier in table_identifiers:
                try:
                    table = tables_input.get_table(table_identifier)
                    tables_to_merge.append(table)
                except KeyError:
                    import warnings

                    warnings.warn(
                        f"Table '{table_identifier}' not found in collection and will be skipped."
                    )

        if not tables_to_merge:
            self._ports_out["OutputTable"].data = None
            return None

        if len(tables_to_merge) == 1:
            # Only one table, return it as-is
            result = tables_to_merge[0]
            self._ports_out["OutputTable"].data = result
            return result

        # Determine the join key
        join_key = self.join_key
        if join_key is None:
            join_key = self._find_join_key(tables_to_merge)
            if join_key is None:
                import warnings

                warnings.warn("No common column found across all tables. Cannot merge.")
                self._ports_out["OutputTable"].data = None
                return None

        # Filter tables that have the join key and resolve column names
        valid_tables: list[TableData] = []
        for table in tables_to_merge:
            resolved_col = self._resolve_column_name(table, join_key)
            if resolved_col is not None:
                # Store the resolved column name with the table for later use
                table._merge_join_key = resolved_col
                valid_tables.append(table)
            else:
                import warnings

                warnings.warn(
                    f"Table '{table.name}' does not contain column '{join_key}' and will be skipped."
                )

        if not valid_tables:
            import warnings

            warnings.warn("No tables contain the specified join key. Cannot merge.")
            self._ports_out["OutputTable"].data = None
            return None

        if len(valid_tables) == 1:
            # Only one valid table, return it as-is
            result = valid_tables[0]
            self._ports_out["OutputTable"].data = result
            return result

        # Perform sequential merging
        result = valid_tables[0]
        for table in valid_tables[1:]:
            # Use the resolved column names for merging
            left_on = getattr(result, "_merge_join_key", join_key)
            right_on = getattr(table, "_merge_join_key", join_key)

            result = result.merge(
                table,
                how=self.how,
                left_on=left_on,
                right_on=right_on,
                suffixes=self.suffixes,
                sort=self.sort,
                copy=self.copy,
                indicator=self.indicator,
                validate=self.validate,
                ignore_index=self.ignore_index,
            )

        # Clean up the temporary attribute
        for table in valid_tables:
            if hasattr(table, "_merge_join_key"):
                delattr(table, "_merge_join_key")

        self._ports_out["OutputTable"].data = result
        return result
